import csv
import glob
import os
from collections import defaultdict

def find_highest_avg_speedup_by_seq_length():
    print("=== 每个 total_seq_length 下最高的平均 Speedup ===\n")
    speedup_files = glob.glob("timings_all_setting_rlhf_*_with_speedup.csv")
    seq_file_avg_speedups = defaultdict(lambda: defaultdict(list))

    for file_path in sorted(speedup_files):
        file_name = os.path.basename(file_path)
        with open(file_path, newline='') as f:
            reader = csv.DictReader(f)
            for row in reader:
                if row.get('speedup') and row['speedup'] != '' and row.get('total_seq_length'):
                    speedup = float(row['speedup'])
                    seq_length = int(row['total_seq_length'])
                    seq_file_avg_speedups[seq_length][file_name].append(speedup)

    # 计算每个 total_seq_length 下每个文件的平均 speedup
    seq_file_avg = defaultdict(dict)
    for seq_length in seq_file_avg_speedups:
        for file_name, speedups in seq_file_avg_speedups[seq_length].items():
            avg_speedup = sum(speedups) / len(speedups)
            seq_file_avg[seq_length][file_name] = avg_speedup

    # 输出每个 total_seq_length 下最高平均 speedup 的文件
    print("total_seq_length | 最高平均speedup | 文件")
    print("-" * 60)
    for seq_length in sorted(seq_file_avg.keys()):
        max_avg_speedup = 0
        best_file = ""
        for file_name, avg_speedup in seq_file_avg[seq_length].items():
            if avg_speedup > max_avg_speedup:
                max_avg_speedup = avg_speedup
                best_file = file_name
        print(f"{seq_length:15d} | {max_avg_speedup:13.4f} | {best_file}")

    # 额外统计：显示所有文件的平均 speedup
    print("\n=== 每个 total_seq_length 下所有文件的平均 Speedup ===\n")
    for seq_length in sorted(seq_file_avg.keys()):
        print(f"total_seq_length {seq_length}:")
        print("-" * 40)
        sorted_files = sorted(seq_file_avg[seq_length].items(), key=lambda x: x[1], reverse=True)
        for file_name, avg_speedup in sorted_files:
            data_count = len(seq_file_avg_speedups[seq_length][file_name])
            print(f"  {file_name}: 平均 {avg_speedup:.4f}, 数据点 {data_count}")
        print()

if __name__ == "__main__":
    find_highest_avg_speedup_by_seq_length()